
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks
import os

import cv2
import numpy as np
from sklearn.model_selection import train_test_split

DATADIR = 'C:\DL_DATA\DATA_SETS\Gesture_Recognition'

IMG_W = 128
IMG_H = 128

ALPHA = 1e-3
BATCH_SIZE = 8
EPOCHS = 20

# data sets
def readData(path):
    x_image = []
    y_label = []
    for i, j in enumerate(os.listdir(path)):
        image_name = os.path.join(path, j)
        for image in os.listdir(image_name):
            image = os.path.join(image_name, image)
            image = cv2.imread(image)
            image = cv2.resize(image, (IMG_W, IMG_H)) / 255
            x_image.append(image)
            y_label.append(i)
    return np.array(x_image), np.array(y_label)

x_image, y_label = readData(DATADIR)

x_train, x_test, y_train, y_test = train_test_split(x_image, y_label, train_size=0.8, shuffle=True)
x_test, x_val, y_test, y_val = train_test_split(x_test, y_test, train_size=0.5, shuffle=True)

def get_model(n_cls, input_shape=(224, 224, 3)):
    inputs = keras.Input(shape=input_shape)
    base_model = keras.applications.MobileNet(
        input_shape=input_shape,
        alpha=1.0,
        depth_multiplier=1,
        dropout=1e-3,
        include_top=False,
        weights='imagenet',
        pooling='avg',
    )
    base_model.trainable = False
    x = base_model(inputs)
    customer_model = keras.layers.Dense(n_cls, activation='softmax')
    x = customer_model(x)
    model = keras.Model(inputs, x)
    return model


if '__main__' == __name__:

    model = get_model(3, (IMG_W, IMG_H, 3))
    model.summary()

    model.compile(
        loss=losses.sparse_categorical_crossentropy,
        optimizer=optimizers.Adam(learning_rate=ALPHA),
        metrics=[metrics.sparse_categorical_accuracy]
    )

    model.fit(x_train, y_train, BATCH_SIZE, EPOCHS,
              validation_data=(x_val, y_val))

    print('Testing...')
    model.evaluate(x_test, y_test, BATCH_SIZE)


    import matplotlib.pyplot as plt
    import random

    pred = np.argmax(model(x_test), 1)

    for i in range(9):
        plt.subplot(3, 3, i + 1)
        r = random.randint(0, len(x_test) - 1)
        img = x_test[r]
        B, G, R = cv2.split(img)
        img2 = cv2.merge([R, G, B])
        if pred[r] == y_test[r]:
            plt.imshow(img2)
            if pred[r] == 0:
                plt.title('left', color="black" )
            elif pred[r] == 1:
                plt.title('right', color="black")
            else:
                plt.title('stop', color="black")
        else:
            plt.imshow(img2)
            if pred[r] == 0:
                plt.title('left', color="red" )
            elif pred[r] == 1:
                plt.title('right', color="red")
            else:
                plt.title('stop', color="red")
        plt.axis('off')
    plt.show()
